// BaseDatasetWrapper.js
const fs = require('fs');
const path = require('path');
const { DatasetListDB } = require('../dbs/dbDatasetList');

class BaseDatasetWrapper {
    constructor(app) {
        // console.log("BaseDatasetWrapper", new.target);
        if (new.target === BaseDatasetWrapper) {
            throw new TypeError('Cannot construct BaseDatasetWrapper instances directly');
        }
        this.app = app;
        this.userDataPath = app.userDataPath;
    }

    // 从目录解析对应的数据集信息
    async parse(dataset_info) {
        throw new Error("Handler method 'parse' must be implemented by subclass.");
    }

    // 导出数据集（导出标注结果和列表）
    async export(dataset_info, subsets_info) {
        throw new Error("Handler method 'export' must be implemented by subclass.");
    }

    // 转换数据集
    async convert(toFormat, dataset_info) {
        throw new Error("Handler method 'convert' must be implemented by subclass.");
    }

    // 保存每一个数据集的image-label list文件
    async saveMatchingPairs(datasetId, matchingPairs) {
        // 创建listdb数据库
        const dbDir = path.join(this.userDataPath, 'datasets', datasetId.toString());
        if (!fs.existsSync(dbDir)) {
            fs.mkdirSync(dbDir, { recursive: true });
        }
        if (!this.app.list_dbs[datasetId]) {
            this.app.list_dbs[datasetId] = new DatasetListDB(this.app, datasetId);
        }

        // 往listdb中写入配对的数据
        await this.app.list_dbs[datasetId].pairs.insertPairs(matchingPairs);
    }

    // 根据dataset的图像目录，创建subset的记录
    async createSubsets(datasetInfo, imageDirs) {
        // 创建listdb数据库
        const datasetId = datasetInfo.dataset_id;
        const dbDir = path.join(this.userDataPath, 'datasets', datasetId.toString());
        if (!fs.existsSync(dbDir)) {
            fs.mkdirSync(dbDir, { recursive: true });
        }
        if (!this.app.list_dbs[datasetId]) {
            this.app.list_dbs[datasetId] = new DatasetListDB(this.app, datasetId);
        }

        // 遍历imageDirs的记录，并作为dir类型的subset插入到数据库中
        const datasetDir = datasetInfo.dataset_dir;
        const workDir = path.join(datasetDir, 'workspace', datasetId);
        let subsetInfo = null;
        for (let i = 0; i < imageDirs.length; i++) {
            let image_dir = imageDirs[i];
            let imageRelativeDir = path.relative(datasetDir, image_dir);
            let labelRelativeDir = path.join(workDir, imageRelativeDir);
            labelRelativeDir = path.relative(datasetDir, labelRelativeDir);
            subsetInfo = {
                dataset_id: datasetId,
                subset_type: 'directory',
                image_dir: imageRelativeDir,
                label_dir: labelRelativeDir,
                label_list_path: ''
            };
            // console.log(`subsetInfo =`, subsetInfo);
            await this.app.list_dbs[datasetId].subsets.insert(subsetInfo);
        }

        // 插入三条记录，训练集、验证集、测试集，默认是空
        // 如果只有一个image目录，则直接给训练集，验证集、测试集都指向同一个目录
        const typelist = ['train', 'val', 'test'];
        let default_image_dir = imageDirs.length === 1 ? subsetInfo.image_dir : '';
        let default_label_dir = imageDirs.length === 1 ? subsetInfo.label_dir : '';
        for (let i = 0; i < typelist.length; i++) {
            subsetInfo = {
                dataset_id: datasetId,
                subset_type: typelist[i],
                image_dir: default_image_dir,
                label_dir: default_label_dir,
                label_list_path: ''
            };
            // console.log(`subsetInfo =`, subsetInfo);
            await this.app.list_dbs[datasetId].subsets.insert(subsetInfo);
        }
    }

    // 保存categories信息
    async saveCategories(datasetId, categories) {
        // console.log('categories = ', categories);
        try {
            for (let id in categories) {
                let category = categories[id];
                let categoryInfo = {
                    category_id: category.id,
                    category_name: category.name,
                    supercategory: category.supercategory,
                    keypoints: JSON.stringify(category.keypoints),
                    skeleton: JSON.stringify(category.skeleton)
                };
                // console.log('category = ', category);
                await this.app.db.categories.insertCategory(datasetId, categoryInfo, '');
            }
            return true;
        } catch (err) {
            console.error('Save Categories Failed! Error = ', err);
            return false;
        }
    }

    async getJsonConfig(format, dataset_info, subsets_info) {
        let jsonConfig = {};

        // 基础信息
        let formatEx = format.toUpperCase();
        let dataset_id = dataset_info.id;
        let dataset_dir = dataset_info.dataset_dir;
        jsonConfig['dataset_type'] = formatEx;
        let categories = await this.app.db.categories.getCategories(dataset_id);
        categories = categories.map((category) => {
            return category.category_name;
        });
        jsonConfig['classes'] = categories;

        // TrainDataset:
        jsonConfig['TrainDataset'] = {};
        let train_subset = subsets_info.find((subset) => subset.subset_type === 'train');
        jsonConfig['TrainDataset']['dataset_dir'] = dataset_dir;
        jsonConfig['TrainDataset']['image_dir'] = train_subset.image_dir;
        if (formatEx !== 'COCO') {
            let label_dir = train_subset.label_dir;
            if (formatEx !== dataset_info.label_format.toUpperCase()) {
                // 数据集格式和转换格式不一致， 需要替换转换后的路径
                let oldDir = path.join('workspace', dataset_id);
                let newDir = path.join('workspace', dataset_id, `converted_${format.toLowerCase()}`);
                label_dir = label_dir.replace(oldDir, newDir);
            }
            jsonConfig['TrainDataset']['anno_dir'] = label_dir;
        }
        jsonConfig['TrainDataset']['anno_path'] = train_subset.label_list_path;

        // EvalDataset:
        jsonConfig['EvalDataset'] = {};
        let eval_subset = subsets_info.find((subset) => subset.subset_type === 'val');
        jsonConfig['EvalDataset']['dataset_dir'] = dataset_dir;
        jsonConfig['EvalDataset']['image_dir'] = eval_subset.image_dir;
        if (formatEx !== 'COCO') {
            let label_dir = eval_subset.label_dir;
            if (formatEx !== dataset_info.label_format.toUpperCase()) {
                // 数据集格式和转换格式不一致， 需要替换转换后的路径
                let oldDir = path.join('workspace', dataset_id);
                let newDir = path.join('workspace', dataset_id, `converted_${format.toLowerCase()}`);
                label_dir = label_dir.replace(oldDir, newDir);
            }
            jsonConfig['EvalDataset']['anno_dir'] = label_dir;
        }
        jsonConfig['EvalDataset']['anno_path'] = eval_subset.label_list_path;

        // TestDataset:
        jsonConfig['TestDataset'] = {};
        let test_subset = subsets_info.find((subset) => subset.subset_type === 'test');
        jsonConfig['TestDataset']['dataset_dir'] = dataset_dir;
        jsonConfig['TestDataset']['image_dir'] = test_subset.image_dir;
        if (formatEx !== 'COCO') {
            let label_dir = test_subset.label_dir;
            if (formatEx !== dataset_info.label_format.toUpperCase()) {
                // 数据集格式和转换格式不一致， 需要替换转换后的路径
                let oldDir = path.join('workspace', dataset_id);
                let newDir = path.join('workspace', dataset_id, `converted_${format.toLowerCase()}`);
                label_dir = label_dir.replace(oldDir, newDir);
            }
            jsonConfig['TestDataset']['anno_dir'] = label_dir;
        }
        jsonConfig['TestDataset']['anno_path'] = test_subset.label_list_path;

        return jsonConfig;
    }

    // 生成训练使用的配置文件
    // dataset_type: COCO
    // classes: [balloon]
    // TrainDataset:
    //     dataset_dir: /home/lixiang/data/balloon/
    //     image_dir: train
    //     anno_path: train/annotation_coco.json
    // EvalDataset:
    //     dataset_dir: /home/lixiang/data/balloon/
    //     image_dir: val
    //     anno_path: val/annotation_coco.json
    // TestDataset:
    //     dataset_dir: /home/lixiang/data/balloon/
    //     image_dir: val
    //     anno_path: val/annotation_coco.json
    async getYmlConfig(format, dataset_info, subsets_info) {
        // console.log(`getConfig format = `, format);
        // console.log(`getConfig dataset_info = `, dataset_info);
        // console.log(`getConfig subsets_info = `, subsets_info);

        // 基础信息
        let formatEx = format.toUpperCase();
        let dataset_id = dataset_info.id;
        let dataset_dir = dataset_info.dataset_dir;
        let config = `dataset_type: ${formatEx}\n`;
        let categories = await this.app.db.categories.getCategories(dataset_id);
        categories = categories.map((category) => {
            return category.category_name;
        });
        config += `classes: [${categories.join(',')}]\n`;

        // TrainDataset:
        config += `TrainDataset:\n`;
        let train_subset = subsets_info.find((subset) => subset.subset_type === 'train');
        config += `    dataset_dir: ${dataset_dir}\n`;
        config += `    image_dir: ${train_subset.image_dir}\n`;
        if (formatEx !== 'COCO') {
            let label_dir = train_subset.label_dir;
            if (formatEx !== dataset_info.label_format.toUpperCase()) {
                // 数据集格式和转换格式不一致， 需要替换转换后的路径
                let oldDir = path.join('workspace', dataset_id);
                let newDir = path.join('workspace', dataset_id, `converted_${format.toLowerCase()}`);
                label_dir = label_dir.replace(oldDir, newDir);
            }
            config += `    anno_dir: ${label_dir}\n`;
        }
        config += `    anno_path: ${train_subset.label_list_path}\n`;

        // EvalDataset:
        config += `EvalDataset:\n`;
        let eval_subset = subsets_info.find((subset) => subset.subset_type === 'val');
        config += `    dataset_dir: ${dataset_dir}\n`;
        config += `    image_dir: ${eval_subset.image_dir}\n`;
        if (formatEx !== 'COCO') {
            let label_dir = eval_subset.label_dir;
            if (formatEx !== dataset_info.label_format.toUpperCase()) {
                // 数据集格式和转换格式不一致， 需要替换转换后的路径
                let oldDir = path.join('workspace', dataset_id);
                let newDir = path.join('workspace', dataset_id, `converted_${format.toLowerCase()}`);
                label_dir = label_dir.replace(oldDir, newDir);
            }
            config += `    anno_dir: ${label_dir}\n`;
        }
        config += `    anno_path: ${eval_subset.label_list_path}\n`;

        // TestDataset:
        config += `TestDataset:\n`;
        let test_subset = subsets_info.find((subset) => subset.subset_type === 'test');
        config += `    dataset_dir: ${dataset_dir}\n`;
        config += `    image_dir: ${test_subset.image_dir}\n`;
        if (formatEx !== 'COCO') {
            let label_dir = test_subset.label_dir;
            if (formatEx !== dataset_info.label_format.toUpperCase()) {
                // 数据集格式和转换格式不一致， 需要替换转换后的路径
                let oldDir = path.join('workspace', dataset_id);
                let newDir = path.join('workspace', dataset_id, `converted_${format.toLowerCase()}`);
                label_dir = label_dir.replace(oldDir, newDir);
            }
            config += `    anno_dir: ${label_dir}\n`;
        }
        config += `    anno_path: ${test_subset.label_list_path}\n`;

        return config;
    }

    // 为了转换格式（复用保存标注结果代码）
    wrapLabelData(dataset_info, pairs_info, label_data) {
        let result = {
            dataset_id: dataset_info.id,
            dataset_type: dataset_info.dataset_type,
            label_format: dataset_info.label_format,
            pairs_id: pairs_info.id,
            image_url: pairs_info.image_path,
            label_url: pairs_info.label_path,
            labels: []
        };
        for (let i = 0; i < label_data.length; i++) {
            // console.log(label_data[i]);
            result.labels.push({
                id: i,
                data: label_data[i]
            });
        }
        // console.log(dataset_info);
        // console.log(pairs_info);
        // console.log(`wrapLabelData result = `, result);
        return result;
    }

    // 获取转换后对应的文件路径
    static getConvertedPath(toFormat, dataset_info, pair) {
        // 格式后缀名
        let format = toFormat.toLowerCase();
        let extname = '.json';
        if (format == 'voc') {
            extname = '.xml';
        } else if (format == 'yolo') {
            extname = '.txt';
        }

        // 数据集相关的属性
        const datasetId = dataset_info.id;
        const datasetDir = dataset_info.dataset_dir;
        const workDir = path.join(datasetDir, 'workspace', `${datasetId}`);
        const ConvertDir = path.join(workDir, `converted_${format}`);

        // 输出路径（按照规则输出）
        let outputPath = path.join(ConvertDir, path.relative(workDir, pair.label_path));
        const basename = path.basename(outputPath, path.extname(outputPath));
        outputPath = path.join(path.dirname(outputPath), basename + extname);

        // 创建目录
        const OutDir = path.dirname(outputPath);
        if (!fs.existsSync(OutDir)) {
            fs.mkdirSync(OutDir, { recursive: true });
        }

        return outputPath;
    }

    // 分析给定的目录标注格式
    static async analyzeFormat(app, dataset_info) {
        // console.log('analyzeFormat dataset_info = ', dataset_info);

        // 根据dataset_supertype和dataset_type确定对应的解析文件所在目录
        const supertype = dataset_info.dataset_supertype;
        const type = dataset_info.dataset_type;
        const parseFilesDir = path.join(__dirname, supertype, type);
        if (fs.existsSync(parseFilesDir) == false) {
            const msg = `Can not find parse files for [${supertype}, ${type}]`;
            console.error(msg);
            return { code: -1, msg: msg };
        }

        // 根据标注文件的类型，调用对应的解析文件
        const wrapperFile = 'Formats.js';
        const wrapperClass = require(path.join(parseFilesDir, wrapperFile));
        const wrapper = new wrapperClass(app);

        // 动态调用解析函数
        let result = await wrapper.parse({
            dataset_dir: dataset_info.dataset_dir
        });
        return result;
    }

    static async invokeParse(app, dataset_info) {
        console.log('invokeParse dataset_info = ', dataset_info);

        // 根据dataset_supertype和dataset_type确定对应的解析文件所在目录
        const supertype = dataset_info.dataset_supertype;
        const type = dataset_info.dataset_type;
        const parseFilesDir = path.join(__dirname, supertype, type);
        if (fs.existsSync(parseFilesDir) == false) {
            const msg = `Can not find parse files for [${supertype}, ${type}]`;
            console.error(msg);
            return { code: -1, msg: msg };
        }

        // 根据标注文件的类型，调用对应的解析文件
        const wrapperFile = dataset_info.label_format.toUpperCase() + 'DatasetWrapper.js';
        const wrapperClass = require(path.join(parseFilesDir, wrapperFile));
        const wrapper = new wrapperClass(app);

        // 动态调用解析函数
        let result = await wrapper.parse({
            dataset_id: dataset_info.id,
            dataset_name: path.basename(dataset_info.dataset_dir),
            dataset_type: dataset_info.dataset_type,
            dataset_dir: dataset_info.dataset_dir,
            label_format: dataset_info.label_format,
            label_selections: dataset_info.label_selections
        });

        // 如果解析成功，则直接插入到数据库中
        if (result.code == 0) {
            // 如果解析成功，则在数据库中插入一条记录
            let totalSamples = result.msg.totalSamples;
            let labeledSamples = result.msg.labeledSamples;
            // id, datasetName, datasetType, datasetDir, labelFormat, totalSamples, labeledSamples
            const lastId = await app.db.datasets.insertDataset({
                id: dataset_info.id,
                datasetName: path.basename(dataset_info.dataset_dir),
                datasetSuperType: dataset_info.dataset_supertype,
                datasetType: dataset_info.dataset_type,
                datasetDir: dataset_info.dataset_dir,
                labelFormat: dataset_info.label_format,
                totalSamples: totalSamples,
                labeledSamples: labeledSamples
            });
            // console.log(`Dataset ${lastId} Created OK! dataset_info = `);
            // console.log(dataset_info);
            result = { code: 0, msg: lastId };
        } else {
            result = { code: -1, msg: 'Parse dataset directory failed!' };
        }

        // 异步返回结果
        app &&
            app.queue.enqueue(
                `dataset_parse:${dataset_info.id}`,
                {
                    event: 'parse_end',
                    data: JSON.stringify(result)
                },
                false
            );

        return result;
    }

    // 定义共同的解析的回调函数，传递解析进度返回到调用方
    parseCallback(datasetId, event, data) {
        this.app &&
            this.app.queue.enqueue(
                `dataset_parse:${datasetId}`,
                {
                    event: event,
                    data: data
                },
                false
            );
    }

    static async invokeConvert(app, to_format, dataset_info) {
        // console.log('invokeConvert dataset_info = ', dataset_info);
        const wrapperFile = dataset_info.label_format.toUpperCase() + 'DatasetWrapper.js';
        const wrapperClass = require(path.join(__dirname, wrapperFile));
        const wrapper = new wrapperClass(app);

        // 动态调用转换函数
        return await wrapper.convert(to_format, dataset_info);
    }

    static async invokeExport(app, train_cfg, dataset_info, subsets_info) {
        // console.log('invokeExport dataset_info = ', dataset_info);
        // 获取外部传入的配置信息
        let toFormat = train_cfg && train_cfg.dataset_type ? train_cfg.dataset_type : 'coco';

        // 先做好格式转换，然后效用对应的导出函数
        let wrapperFile = dataset_info.label_format.toUpperCase() + 'DatasetWrapper.js';
        let wrapperClass = require(path.join(__dirname, wrapperFile));
        let wrapper = new wrapperClass(app);
        await wrapper.convert(toFormat, dataset_info);

        // 动态调用导出函数
        wrapperFile = toFormat.toUpperCase() + 'DatasetWrapper.js';
        wrapperClass = require(path.join(__dirname, wrapperFile));
        wrapper = new wrapperClass(app);
        await wrapper.export(toFormat, dataset_info, subsets_info);

        // 输出数据集配置文件
        // let config = await wrapper.getYmlConfig(toFormat, dataset_info, subsets_info);
        let config = await wrapper.getJsonConfig(toFormat, dataset_info, subsets_info);
        return config;
    }
}

module.exports = BaseDatasetWrapper;
